# problem 1

import matplotlib.pyplot as plt
import numpy as np
from PIL import Image

def compress_image(img_path, num_components):
    # Load color image
    img = Image.open(img_path)
    img_array = np.array(img)

    # Separate color channels (R, G, B)
    red_channel = img_array[:,:,0]
    green_channel = img_array[:,:,1]
    blue_channel = img_array[:,:,2]

    # Perform SVD for each channel
    U_red, sigma_red, V_red = np.linalg.svd(red_channel, full_matrices=False)
    U_green, sigma_green, V_green = np.linalg.svd(green_channel, full_matrices=False)
    U_blue, sigma_blue, V_blue = np.linalg.svd(blue_channel, full_matrices=False)

    # Keep only the specified number of components
    U_red = U_red[:, :num_components]
    sigma_red = np.diag(sigma_red[:num_components])
    V_red = V_red[:num_components, :]

    U_green = U_green[:, :num_components]
    sigma_green = np.diag(sigma_green[:num_components])
    V_green = V_green[:num_components, :]

    U_blue = U_blue[:, :num_components]
    sigma_blue = np.diag(sigma_blue[:num_components])
    V_blue = V_blue[:num_components, :]

    # Reconstruct compressed channels
    compressed_red_channel = U_red @ sigma_red @ V_red
    compressed_green_channel = U_green @ sigma_green @ V_green
    compressed_blue_channel = U_blue @ sigma_blue @ V_blue

    # Stack the compressed channels to get the compressed color image
    compressed_image = np.stack(
        [compressed_red_channel, compressed_green_channel, compressed_blue_channel],
        axis=-1
    )

    # Display the original and compressed images
    plt.figure(figsize=(12, 6))

    # Original Image
    plt.subplot(1, 2, 1)
    plt.imshow(img)
    plt.title('Original Image')

    # Compressed Image
    plt.subplot(1, 2, 2)
    plt.imshow(compressed_image.astype(np.uint8))
    plt.title(f'Compressed Image ({num_components} components)')

    plt.show()

# Example usage
compress_image('/content/einstein-color-image.jpg', num_components=50)

# problem 2

# rayleight itertion method
def rayleigh_iteration(A, a, num_iterations=1000, tol=1e-10):
    identity = np.eye(A.shape[0])

    for _ in range(num_iterations):
        a_norm = np.linalg.norm(a)
        a = a / a_norm
        sigma = np.dot(np.dot(a.T, A), a)
        matrix_inverse = np.linalg.inv(A - sigma * identity)
        a = np.dot(matrix_inverse, a)

    a_norm = np.linalg.norm(a)
    a = a / a_norm
    sigma = np.dot(np.dot(a.T, A), a)

    return sigma, a

# power_iteration method
def power_iteration(A, num_iterations=1000, tol=1e-10):
    n = A.shape[0]
    a = np.random.rand(n)

    for _ in range(num_iterations):
        a_next = np.dot(A, a)
        a_next_norm = np.linalg.norm(a_next)
        a = a_next / a_next_norm

    return a


def svd(A, tol=1e-10, num_iterations=1000):
    AAt = np.dot(A, A.T)
    AtA = np.dot(A.T, A)

    eigvectors = power_iteration(AAt, num_iterations=num_iterations, tol=tol)

    eig_val, eig_vec = rayleigh_iteration(AAt, eigvectors,
                                                 num_iterations=num_iterations, tol=tol)

    S = np.sqrt(np.abs(eig_val))

    U = eig_vec
    V = np.dot(A.T, U) / S

    return U, S, V


np.random.seed(100)
A = np.random.random((3, 3))
U, S, Vt = svd(A)

print("U :", U)
print("S :", S)
print("Vt :", Vt)